/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Definitions of the full_sed_grid functions.

#include "full_sed_grid.h"
#include "grid-fits.h"
#include <CCfits/CCfits>
#include "blitz-fits.h"
#include <vector>
#include "density_generator.h"
#include "biniostream.h"
#include "fits-utilities.h"
#include "constants.h"
#include "grid.h"
#include "arepo_grid.h"
#include "emission.h"
#include "arepo_emission.h"

#ifdef WITH_AREPO
extern "C" {
#include "allvars.h"
#include "proto.h"
}
#endif

using namespace CCfits;
using namespace std;
using namespace blitz;

template<typename T>
mcrx::full_sed_grid<T>::full_sed_grid (boost::shared_ptr<T_grid_impl> g,
				       CCfits::ExtHDU& structure_hdu,
				       CCfits::ExtHDU& data_hdu,
				       const density_generator& gen,
				       int n_lambda) :
  polychromatic_grid<T> (g, data_hdu, gen, n_lambda)
{
  // this constructor only has to read units.

  structure_hdu.readKey("lengthunit", units_ ["length"]);
  Column& c_m_g = data_hdu.column("mass_gas");
  units_ ["mass"] = c_m_g.unit();
}


template<typename T>
mcrx::full_sed_grid<T>::full_sed_grid (boost::shared_ptr<T_grid_impl> g,
				       const T_unit_map& units):
  polychromatic_grid<T> (g)
{
  // this constructor only has to copy units.
  units_.insert(units.begin(), units.end());
  // however, this polychromatic_grid constructor doesn't call
  // make_blocks, so we need to do that here. The actual cells are
  // already set up in the grid, since we used a factory.
  this->make_blocks();
}


template<typename T>
mcrx::full_sed_grid<T>::full_sed_grid (boost::shared_ptr<T_grid_impl> g,
				       const density_generator& gen,
				       int n_lambda) :
  polychromatic_grid<T> (g)
{
#ifdef WITH_AREPO
  const int nc = this->n_cells ();
  assert(nc==N_gas);
  cout << "Creating cell data from Arepo data for " << nc << " cells. " << endl;

  // First copy units from the arepo units
  units_["length"]=arepo::arepo_units.get("length");
  units_["time"]=arepo::arepo_units.get("time");
  units_["mass"]=arepo::arepo_units.get("mass");

  // to convert arepo density to our units. (convert call kind of
  // unnecessary since we know the units are the same from above, but
  // done for clarity)
  const T_float convert_density = 
    arepo::mcon/(arepo::lcon*arepo::lcon*arepo::lcon)*
    units::convert(arepo::arepo_units.get("mass")+"/"+
		   arepo::arepo_units.get("length")+"^3",
		   units_.get("mass")+"/"+units_.get("length")+"^3");

  // to convert arepo velocity to our units.
  const T_float convert_velocity = 
    arepo::lcon/arepo::tcon*
    units::convert(arepo::arepo_units.get("length")+"/"+
		   arepo::arepo_units.get("time"),
		   units_.get("length")+"/"+units_.get("time"));

  int i = 0;
  bool negz=false;
  for (iterator c = this->begin (), e = this->end(); c != e; ++c, ++i) {
    // make an absorber with the appropriate densities. We only know
    // density of metals and gas so we ask the density generator to
    // give us a density vector, which we put into our big block.
    assert(c.volume()>0);
    const T_float rho_gas = SphP[i].Density*convert_density;

    // for some reason, Arepo gets tiny negative metallicities. If
    // this happens, set it to zero and print a warning.

    if(P[i].Metallicity<0) {
      negz=true;
      assert(P[i].Metallicity>-1e-6);
      P[i].Metallicity=0;
    }

    const T_float rho_met = rho_gas*P[i].Metallicity;
    // velocity has no factors of h
    vec3d vel(SphP[i].Momentum[0], SphP[i].Momentum[1], SphP[i].Momentum[2]);
    vel *= convert_velocity/(SphP[i].Density*SphP[i].Volume);

    const T_densities rho (gen.make_density_vector(rho_gas, rho_met));
    ASSERT_ALL(rho>=0);
    DEBUG(3,cout << "Cell " << i << " rho_gas=" << rho_gas << ", rho_met=" << rho_met << " rho=" << rho(0) << endl;);
    T_lambda intens (n_lambda);
    intens=0;
    T_data d(rho, vel, intens);

    // if the cell has no data member, default-construct one before assigning
    if(!c->data()) {
      c->set_data(new typename T_grid_impl::T_data() );
    }
    c->data()->get_absorber() = d;
  }
  if(negz)
    cerr << "Warning: Negative metallicities found in Arepo grid.\n";

  // and now consolidate into continuous blocks
  this->make_blocks();

#else
  cerr << "Not compiled with Arepo support" << endl;
  exit(1);
#endif
}


template<typename T>
std::pair<ExtHDU*, ExtHDU*>
mcrx::full_sed_grid<T>::create_hdus(CCfits::FITS& file, 
				    const TinyVector<int,2>& size,
				    const std::string& i_name,
				    const std::string& d_name,
				    const string& luminosityunit,
				    bool do_compress) const
{
  std::vector<long> naxes,znaxes;
  naxes.push_back(size[0]);
  naxes.push_back(size[1]);
  znaxes=naxes;
  znaxes[0]=100; // compress 100 cells as one
  
  // check if intensity HDU already exists
  ExtHDU* i_hdu;    
  try {
    i_hdu = &open_HDU (file, i_name);
  }
  catch (CCfits::FITS::NoSuchHDU&) {
    cout << "Creating " << i_name << " HDU with size " << size << endl;
    if (do_compress) {
      // use the cfitsio tile compression on these images
      file.setCompressionType(GZIP_1);
      file.setTileDimensions(znaxes);
    }
    
    i_hdu = file.addImage (i_name, FLOAT_IMG, naxes);
    file.setCompressionType(0);
    
    i_hdu->writeComment("This HDU contains the mean radiation intensity in the grid cells due to primary (stellar) emission. The first dimension is grid cell, the second wavelength.");

    // area unit is hard-coded to be m^2 so we take whatever we're fed
    // and add /m^2
    const string intunit = luminosityunit+"/m^2";
    i_hdu->addKey("imunit", intunit, "Mean intensity in units of "+intunit);
  }
  
  // check if deposition HDU already exists
  ExtHDU* d_hdu;    
  try {
    d_hdu = &open_HDU (file, d_name);
  }
  catch (CCfits::FITS::NoSuchHDU&) {
    cout << "Creating " << d_name << " HDU with size " << size << endl;
    if (do_compress) {
      // use the cfitsio tile compression on these images
      file.setCompressionType(GZIP_1);
      file.setTileDimensions(znaxes);
    }
    
    d_hdu = file.addImage (d_name, FLOAT_IMG, naxes);
    file.setCompressionType(0);
    
    d_hdu->writeComment("This HDU contains the energy deposition rate in the grid cells. The first dimension is grid cell, the second wavelength.");
  }

  return make_pair(i_hdu, d_hdu);
}



/** Loads the radiation intensity in the grid cells from a FITS
    HDU. If we're running on several tasks, only loads the subset of
    cells owned by this task. */
template<typename T>
void mcrx::full_sed_grid<T>::load_intensity (const CCfits::ExtHDU& i_hdu)
{
  cout << "Loading grid cell intensities from file" << endl;

  //figure out where "our" part of the cell array starts.
  const int nc = this->n_cells();
  const int start_cell = mpi_calculate_offsets(nc);

  // Ideally, we would read directly into the big intensity array, but
  // because the values have to be converted using the cell volume, we
  // have to look anyway, so we just use a temporary array for
  // simplicity.

  // Indices are (scatterer, lambda).
  array_2 intensity;
  read(i_hdu, intensity, nc, start_cell);
  assert(intensity.extent(firstDim)==nc);

  intensity = pow(10, intensity);

  float normalization = 1;
  T_float area_factor = 1;
  try {
    const_cast<ExtHDU&> (i_hdu).readKey ("normalization",
						 normalization);
  }
  catch (...) {}
  try {
    const_cast<ExtHDU&> (i_hdu).readKey ("area_factor",
						 area_factor);
  }
  catch (...) {}

  int ii=0;
  for (const_iterator c = this->begin (); c != this->end (); ++c, ++ii) {
    assert (ii < intensity.rows());
    // invert the calculation done in write_intensity here
    const T_lambda temp
      (intensity(ii, Range::all()) * 4*constants::pi*c.volume()*area_factor /
       normalization);
    // set intensity, remember this puts it into the big array.
    c->data()->get_absorber().set_intensity(temp);
  }
}

/** Dumps the radiation intensity to a file. */
template<typename T>
void mcrx::full_sed_grid<T>::write_dump (binofstream& file) const
{
  // We already have the intensities in one large array so this is easy.
  TinyVector<int, 2> extent = this->intensities_.shape();
  DEBUG(1,cout << "Dumping intensity data for grid, " << extent << " bytes" << endl;);
  file << extent [0] << extent [1];
  file.write(reinterpret_cast<const char*> (this->intensities_.dataFirst()),
	     this->intensities_.size()*sizeof (T_float) );
}

/** Loads the absorbed luminosity and radiation intensity in the grid
    cells from a previously written dump file. */
template<typename T>
void mcrx::full_sed_grid<T>::load_dump (binifstream& file) 
{
  // We should already have an appropriately sized large array for the
  // intensities, so loading is easy.
  TinyVector< int, 2> extent;
  file >> extent [0] >>  extent [1];
  DEBUG(1,cout << "Loading intensity data for grid, " << extent << " bytes" << endl;);
  assert(extent[0] == this->intensities_.shape()[0]);
  assert(extent[1] == this->intensities_.shape()[1]);

  file.read(reinterpret_cast<char*> (this->intensities_.dataFirst()),
	    this->intensities_.size()*sizeof (T_float) );
}


// Explicitly intantiate the templates necessary to avoid having to
// include everything. UNfortunately we can't instantiate the typedefs
// we made in full_sed_grid.h...
namespace mcrx {
  template class full_sed_grid<
    adaptive_grid<
      cell_data<
	typename emitter_types<adaptive_grid, 
			       polychromatic_policy, local_random>::T_emitter,
	absorber<array_1> > > >;

  template class full_sed_grid<adaptive_grid<absorber<array_1> > >;

#ifdef WITH_AREPO
  template class full_sed_grid<
    arepo_grid<
      cell_data<
	typename emitter_types<arepo_grid, 
			       polychromatic_policy, local_random>::T_emitter,
	absorber<array_1> > > >;
#endif
}
